package main

import (
	"bufio"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/gocolly/colly/v2"
	"github.com/gocolly/colly/v2/debug"
	"github.com/spf13/viper"
	"github.com/valyala/fastjson"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/rs/zerolog/log"
)

var (
	baseUrl string
	start   int
	end     int
)

var userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.77 Safari/537.36"

var header1 = http.Header{
	"accept":                         []string{"*/*"},
	"accept-encoding":                []string{"gzip, deflate, br"},
	"accept-language":                []string{"zh-CN,zh;q=0.9"},
	"access-control-request-headers": []string{"range"},
	"access-control-request-method":  []string{"GET"},
	"cache-control":                  []string{"no-cache"},
	"origin":                         []string{"https://www.bilibili.com"},
	"pragma":                         []string{"no-cache"},
	"referer":                        []string{"https://www.bilibili.com/"},
	"sec-fetch-dest":                 []string{"empty"},
	"sec-fetch-mode":                 []string{"cors"},
	"sec-fetch-site":                 []string{"cross-site"},
	"user-agent":                     []string{userAgent},
}

var header2 = http.Header{
	"accept":          []string{"*/*"},
	"accept-encoding": []string{"identity"},
	"accept-language": []string{"zh-CN,zh;q=0.9,en;q=0.8"},
	"cache-control":   []string{"no-cache"},
	"origin":          []string{"https://www.bilibili.com"},
	"pragma":          []string{"no-cache"},
	//"referer":         []string{"https://www.bilibili.com/video/BV1MB4y1M7Au"},
	"sec-ch-ua":        []string{`" Not;A Brand"; v = "99", "Google Chrome"; v = "91", "Chromium"; v = "91"`},
	"sec-fetch-dest":   []string{"empty"},
	"sec-fetch-mode":   []string{"cors"},
	"sec-fetch-site":   []string{"cross-site"},
	"user-agent":       []string{userAgent},
	"sec-ch-ua-mobile": []string{"?0"},
}

const (
	Mega          = 1024 * 1024
	RangeStartKey = "RangeStartKey"
	RangeEndKey   = "RangeEndKey"

	RcvBufferKey = "RcvBuffer"

	RequestType    = "RequestType"
	OPTIONSVideo   = "OPTIONS Video"
	OPTIONSAudio   = "OPTIONS Audio"
	GetHTMLRequest = "GETHtml"

	MediaTypeKey   = "MediaTypeKey"
	MediaTypeVideo = 0
	MediaTypeAudio = 1
)

func init() {
	flag.StringVar(&baseUrl, "url", "", "-url xxx.xxx.xxx")
	flag.IntVar(&start, "start", 1, "-start 1")
	flag.IntVar(&end, "end", 1, "-end 5")
}

func main() {
	flag.Parse()

	setupConfig()
	setupLog()

	if len(baseUrl) == 0 {
		return
	}

	for i := start; i <= end; i++ {
		err := getBilibiliMedia(baseUrl, i)
		if err != nil {
			return
		}
	}
}

func getBilibiliMedia(URL string, episodeNo int) error {
	hc := NewHtmlCollector(URL, episodeNo)

	err := hc.Request("GET")
	if err != nil {
		log.Err(err).Msg("hc.Request")
		return err
	}

	return nil
}

func (mc *MediaCollector) onMediaAllReceived(r *colly.Response) error {
	mediaID := getMediaID(mc.BaseURL)
	mediaType := r.Ctx.GetAny(MediaTypeKey).(MediaType)
	var filename string
	if mediaType == MediaTypeVideo {
		filename = fmt.Sprintf("%s_%04d.video", mediaID, mc.episode)
	} else {
		filename = fmt.Sprintf("%s_%04d.audio", mediaID, mc.episode)
	}
	rcvBuf := r.Ctx.GetAny(RcvBufferKey).(*bytes.Buffer)
	err := saveToFile(rcvBuf.Bytes(), filename)
	if err != nil {
		return err
	}
	log.Info().Str("fileName", filename).Msg("save to file")

	if mediaType == MediaTypeAudio {
		filename := fmt.Sprintf("%s_%04d", mediaID, mc.episode)
		mergeVideoAndAudio(filename)
	}
	return nil
}

func mergeVideoAndAudio(filename string) {

	videoFileName := "./mpeg/" + filename + ".video"
	audioFileName := "./mpeg/" + filename + ".audio"
	outputFileName := "./mpeg/" + filename + ".mp4"

	cmd := exec.Command(
		"ffmpeg",
		"-i",
		videoFileName,
		"-i",
		audioFileName,
		"-vcodec",
		"copy",
		"-acodec",
		"copy",
		outputFileName,
	)
	cmd.Stdout = os.Stdout

	err := cmd.Run()
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			if exitErr.Stderr == nil {
				log.Info().Str("Mp4File", outputFileName).Msg("Video Audio file merged success")
				err := os.Remove(videoFileName)
				if err != nil {
					log.Err(err).Send()
					return
				}
				err = os.Remove(audioFileName)
				if err != nil {
					log.Err(err).Send()
					return
				}
				return
			}
		} else {
			log.Err(err).Send()
		}
	}

	err = os.Remove(videoFileName)
	if err != nil {
		log.Err(err).Send()
		return
	}
	err = os.Remove(audioFileName)
	if err != nil {
		log.Err(err).Send()
		return
	}

}

func getMediaID(URL string) string {
	idx := strings.LastIndex(URL, "/")
	return URL[idx+1:]
}

func (mc *MediaCollector) Request(
	c *colly.Collector,
) error {
	ctx := colly.NewContext()
	ctx.Put(RangeStartKey, 0)
	//ctx.Put(RangeEndKey, viper.GetInt("capture.one_request_max")*Mega)
	ctx.Put(RangeEndKey, 256)
	ctx.Put(MediaTypeKey, mc.mediaType)

	var buf []byte
	rcvBuf := bytes.NewBuffer(buf)
	ctx.Put(RcvBufferKey, rcvBuf)

	_, err := mc.requestMedia(c, ctx)
	if err != nil {
		log.Err(err).Send()
	}
	return err
}

func getTotalSize(contentRange string) int {
	idx := strings.LastIndex(contentRange, "/")
	result, _ := strconv.Atoi(contentRange[idx+1:])
	return result
}

func getTargetURLFromHTML(htmlBody []byte) (videoURL []byte, audioURL []byte, err error) {
	reg := regexp.MustCompile(`<script>window.__playinfo__=(.*?)</script>`)
	data := reg.FindSubmatch(htmlBody)
	if len(data) <= 1 {
		return nil, nil, errors.New("can't find url")
	}
	videoUrl, audioUrl, err := getVideoUrl(data[1])
	if err != nil {
		return nil, nil, err
	}
	return videoUrl, audioUrl, nil
}

func (mc *MediaCollector) requestMedia(c *colly.Collector,
	ctx *colly.Context) ([]*http.Cookie, error) {
	cookies := c.Cookies(mc.URL)
	err := mc.c.SetCookies(mc.URL, cookies)
	if err != nil {
		log.Err(err)
	}
	err = mc.c.Request("GET", mc.URL,
		nil, ctx, header2)
	if err != nil {
		log.Err(err)
	}
	return cookies, err
}

func printHeader(headers *http.Header) {
	if e := log.Debug(); e.Enabled() {
		log.Trace().Str("start", "---------------------------------------").
			Send()
		headers.Write(os.Stdout)
		log.Trace().Str("end", "=======================================").
			Send()
	}
}

func getVideoUrl(jsonData []byte) ([]byte, []byte, error) {
	err := fastjson.ValidateBytes(jsonData)
	if err != nil {
		return nil, nil, err
	}
	values, err := fastjson.ParseBytes(jsonData)
	if err != nil {
		return nil, nil, err
	}
	videos := values.GetArray("data", "dash", "video")
	var v, a []byte
	v = videos[0].GetStringBytes("baseUrl")

	audios := values.GetArray("data", "dash", "audio")

	a = audios[0].GetStringBytes("baseUrl")

	return v, a, nil
}

func saveToFile(p []byte, filename string) error {
	filename = "./mpeg/" + filename
	dstFile, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, os.ModePerm)
	if err != nil {
		return err
	}
	bufWriter := bufio.NewWriter(dstFile)
	defer func() {
		bufWriter.Flush()
		dstFile.Close()
	}()

	_, err = bufWriter.Write(p)
	if err != nil {
		return err
	}
	return nil
}

// HtmlCollector is the main Collector
type HtmlCollector struct {
	URL     string
	episode int
	c       *colly.Collector
	vc      *MediaCollector
	ac      *MediaCollector
}

func NewHtmlCollector(URL string, episode int) *HtmlCollector {
	hc := &HtmlCollector{
		URL:     URL,
		episode: episode,
	}

	if e := log.Debug(); e.Enabled() {
		hc.c = colly.NewCollector(
			colly.UserAgent(userAgent),
			colly.AllowURLRevisit(),
			colly.ParseHTTPErrorResponse(),
			colly.Debugger(&debug.LogDebugger{}),
			colly.MaxBodySize(viper.GetInt("capture.max_body_size")*Mega))
	} else {
		hc.c = colly.NewCollector(
			colly.UserAgent(userAgent),
			colly.AllowURLRevisit(),
			colly.ParseHTTPErrorResponse(),
			colly.MaxBodySize(viper.GetInt("capture.max_body_size")*Mega))
	}
	hc.c.SetRequestTimeout(time.Second * viper.GetDuration("capture.timeout"))
	hc.c.OnResponse(hc.OnHtmlCollectorResponse)
	hc.c.OnRequest(hc.OnHtmlCollectorRequest)
	hc.c.OnError(hc.OnHtmlCollectorError)

	hc.vc = NewMediaCollector(hc.c, MediaTypeVideo, episode, URL)

	hc.ac = NewMediaCollector(hc.c, MediaTypeAudio, episode, URL)

	return hc
}

func (hc *HtmlCollector) OnHtmlCollectorError(r *colly.Response, err error) {
	log.Err(err).Msg("OnHtmlCollectorError.OnError")
}

func (hc *HtmlCollector) Request(
	method string,
) error {
	ctx := colly.NewContext()
	ctx.Put(RequestType, GetHTMLRequest)
	realUrl := fmt.Sprintf("%s?p=%d", hc.URL, hc.episode)
	return hc.c.Request(method, realUrl, nil, ctx, header1)
}

func (hc *HtmlCollector) OnHtmlCollectorRequest(r *colly.Request) {
	log.Debug().Str("Method", r.Method).
		Str("URL", r.URL.String()).
		Msg("c.OnRequest")
}

func (hc *HtmlCollector) OnHtmlCollectorResponse(r *colly.Response) {
	switch r.Ctx.Get(RequestType) {
	case GetHTMLRequest:
		log.Debug().Str("RequestType", GetHTMLRequest).
			Int("body size", len(r.Body)).
			Strs("content-type", r.Headers.Values("content-type")).
			Send()

		videoUrl, audioUrl, err := getTargetURLFromHTML(r.Body)
		if err != nil {
			log.Err(err).Msg("getTargetURLFromHTML")
			return
		}
		log.Debug().Str("videoURL", string(videoUrl)).
			Str("audioURL", string(audioUrl)).Msg("c.OnResponse")

		hc.vc.URL = string(videoUrl)
		hc.ac.URL = string(audioUrl)

		r.Ctx.Put(RequestType, OPTIONSVideo)
		err = hc.c.Request("OPTIONS", string(videoUrl),
			nil, r.Ctx, header1)
		if err != nil {
			log.Err(err).Msg("request option Video")
			return
		}

		r.Ctx.Put(RequestType, OPTIONSAudio)
		err = hc.c.Request("OPTIONS", string(audioUrl),
			nil, r.Ctx, header1)
		if err != nil {
			log.Err(err).Msg("request option Audio")
			return
		}
	case OPTIONSVideo:
		log.Debug().Str("RequestType", OPTIONSVideo).
			Send()
		err := hc.vc.Request(hc.c)
		if err != nil {
			log.Err(err).Send()
			return
		}
	case OPTIONSAudio:
		log.Debug().Str("RequestType", OPTIONSAudio).
			Send()
		err := hc.ac.Request(hc.c)
		if err != nil {
			log.Err(err).Send()
			return
		}
	default:
		log.Err(errors.New("default request type")).Send()
		return
	}
}

type MediaType int
type MediaCollector struct {
	URL       string
	BaseURL   string
	mediaType MediaType
	episode   int
	c         *colly.Collector
}

func NewMediaCollector(
	clone *colly.Collector,
	mt MediaType,
	episode int,
	baseURL string,
) *MediaCollector {
	mc := &MediaCollector{}
	mc.BaseURL = baseURL
	mc.episode = episode
	mc.mediaType = mt
	mc.c = clone.Clone()
	mc.c.OnResponse(mc.OnMediaCollectorResponse)
	mc.c.OnRequest(mc.OnMediaCollectorRequest)
	mc.c.OnError(mc.OnMediaCollectorError)
	return mc
}

func (mc *MediaCollector) OnMediaCollectorResponse(r *colly.Response) {
	strContentLength := r.Headers.Get("Content-Length")
	strContentRange := r.Headers.Get("Content-Range")
	log.Debug().
		Str("Content-Length", strContentLength).
		Str("Content-Range", strContentRange).
		Send()

	if len(r.Body) == 0 {
		log.Info().Msg("Receive 0 byte streamMedia")
		return
	}

	rcvBuf := r.Ctx.GetAny(RcvBufferKey).(*bytes.Buffer)
	rcvBuf.Write(r.Body)

	totalSize := getTotalSize(strContentRange)
	rangStart := r.Ctx.GetAny(RangeStartKey).(int)
	rangEnd := r.Ctx.GetAny(RangeEndKey).(int)
	rangStart = rangStart + len(r.Body)
	if rangStart >= totalSize-1 {
		log.Info().Msg("all media received")
		err := mc.onMediaAllReceived(r)
		if err != nil {
			log.Err(err).Send()
		}
		return
	}

	rangEnd = rangStart + viper.GetInt("capture.one_request_max")*Mega
	if rangEnd >= totalSize {
		rangEnd = totalSize - 1
	}
	r.Ctx.Put(RangeStartKey, rangStart)
	r.Ctx.Put(RangeEndKey, rangEnd)
	_, err := mc.requestMedia(
		mc.c,
		r.Ctx,
	)
	if err != nil {
		log.Err(err).Send()
		return
	}
}

func (mc *MediaCollector) OnMediaCollectorRequest(r *colly.Request) {
	rangeStart := r.Ctx.GetAny(RangeStartKey)
	rangeEnd := r.Ctx.GetAny(RangeEndKey)

	r.Headers.Del("range")
	r.Headers.Del("referer")
	ranges := fmt.Sprintf("bytes=%d-%d", rangeStart, rangeEnd)
	r.Headers.Add("range", ranges)
	r.Headers.Add("referer", baseUrl)
	printHeader(r.Headers)
}

func (mc *MediaCollector) OnMediaCollectorError(r *colly.Response, err error) {
	log.Err(err).Msg("mediaCollector.OnError")
}
